#!/usr/bin/python3
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

import numpy
import numpy.linalg
import scipy.sparse
import scipy.sparse.linalg

def maxEigVec(adjmat):
  '''
# Calculate maximum eigenvalue of Symmetrized Adjacent Matrix and its eigenvector
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  adjacent matrix
# output:
#  {eigenvalue, eigenvector}
  '''
  adjmat=adjmat.tocsr();
  symadj=adjmat+adjmat.transpose(); # should be (A^T+A)/2, the 'div 2' is put in return expression.
  egval,egvec=scipy.sparse.linalg.eigs(symadj,k=1,which='LM'); # eigenvalue with largest magnitude
  egvec=numpy.real(egvec[:,0]);
  egvec=egvec*(egvec>0)-egvec*(egvec<0); # For a positive matrix, all eigenvector components of max eigenvalue should have the same sign.
  return {"eigenvalue":egval[0]/2.,"eigenvector":egvec};

def flowMat(adjmat,clusters):
  '''
# Calculate flow matrix of a given clustering
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  adjacent matrix, cluster division (list)
# output:
#  flow matrix
  '''
  dimp=adjmat.shape[0];
  if (dimp!=adjmat.shape[1]):
   print('WARNING: adjacent matrix is not square. Row:'+str(adjmat.shape[0])+' ,Col:'+str(adjmat.shape[1]));
  dimc=len(clusters);
  adjcsr=adjmat.tocsr();
  flowmat=numpy.matrix([[adjcsr[clusters[c1]].transpose()[clusters[c2]].sum()/numpy.sqrt(len(clusters[c1])*len(clusters[c2])) for c2 in range(0,dimc)] for c1 in range(0,dimc)]); # flowmat is dimc x dimc
  return flowmat;

def pip_flowMat(lastflowmat,numvec,clusters):
  '''
# Calculate flow matrix of a given clustering
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  last flow matrix, last paper number vector, cluster division
# output:
#  flow matrix, paper number vector of new clusters
  '''
  dimp=lastflowmat.shape[0];
  if (dimp!=lastflowmat.shape[1]):
   print('WARNING: last flow matrix is not square. Row:'+str(lastflowmat.shape[0])+' ,Col:'+str(lastflowmat.shape[1]));
  dimc=len(clusters);
  arrnumvec=numpy.array(numvec); # use numpy.array to slice
  adjcsr=lastflowmat.tocsr();
  newflowmat=numpy.matrix([[(scipy.sparse.diags(numpy.sqrt(arrnumvec[clusters[c2]]),offsets=0)*((scipy.sparse.diags(numpy.sqrt(arrnumvec[clusters[c1]]),offsets=0)*adjcsr[clusters[c1]]).transpose()[clusters[c2]])).sum()/numpy.sqrt(arrnumvec[clusters[c1]].sum()*arrnumvec[clusters[c2]].sum()) for c2 in range(0,dimc)] for c1 in range(0,dimc)]); # flowmat is dimc x dimc
  newnumvec=[arrnumvec[clusters[c]].sum() for c in range(0,dimc)];
  return newflowmat,newnumvec;

def pip_flowMat_sparse(lastflowmat,numvec,clusters): # Slow!
  '''
# Calculate flow matrix of a given clustering (sparse mode) [slow!]
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  last flow matrix, last paper number vector, cluster division
# output:
#  flow matrix, paper number vector of new clusters
  '''
  dimp=lastflowmat.shape[0];
  if (dimp!=lastflowmat.shape[1]):
   print('WARNING: last flow matrix is not square. Row:'+str(lastflowmat.shape[0])+' ,Col:'+str(lastflowmat.shape[1]));
  dimc=len(clusters);
  arrnumvec=numpy.array(numvec); # use numpy.array to slice
  adjcsr=lastflowmat.tocsr();
  flowsVal=[];
  flowsRow=[];
  flowsCol=[];
  for c1 in range(0,dimc):
   for c2 in range(0,dimc):
    flow=(scipy.sparse.diags(numpy.sqrt(arrnumvec[clusters[c2]]),offsets=0)*((scipy.sparse.diags(numpy.sqrt(arrnumvec[clusters[c1]]),offsets=0)*adjcsr[clusters[c1]]).transpose()[clusters[c2]])).sum()/numpy.sqrt(arrnumvec[clusters[c1]].sum()*arrnumvec[clusters[c2]].sum());
    if (flow!=0.):
     flowsVal.append(flow);
     flowsRow.append(c1);
     flowsCol.append(c2);
  newflowmat=scipy.sparse.coo_matrix((numpy.array(flowsVal), (numpy.array(flowsRow), numpy.array(flowsCol))),shape=(dimc,dimc));
  newnumvec=[arrnumvec[clusters[c]].sum() for c in range(0,dimc)];
  return newflowmat,newnumvec;

def simfMat(flowmat):
  '''
# Calculate forward similarity matrix of a given flow matrix
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  flow matrix
# output:
#  forward similarity matrix
  '''
  flowmat=flowmat.copy();
  for it in range(0,numpy.min(flowmat.shape)):
    flowmat.A[it][it]=0;
  return flowmat*(flowmat.transpose());

def simbMat(flowmat):
  '''
# Calculate backward similarity matrix of a given flow matrix
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  flow matrix
# output:
#  backward similarity matrix
  '''
  flowmat=flowmat.copy();
  for it in range(0,numpy.min(flowmat.shape)):
    flowmat.A[it][it]=0;
  return (flowmat.transpose())*flowmat;

def sim1Mat(flowmat):
  '''
# Calculate similarity matrix of a given flow matrix
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  flow matrix
# output:
#  similarity matrix
  '''
  return (simfMat(flowmat)+simbMat(flowmat))/2.;

def tflfVec(flowmat):
  '''
# Calculate forward total flow vector of a given flow matrix
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  flow matrix
# output:
#  forward total flow vector
  '''
  return numpy.array((numpy.matrix(numpy.ones(flowmat.shape[1]))*(flowmat.transpose())).tolist());

def tflbVec(flowmat):
  '''
# Calculate backward total flow vector of a given flow matrix
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  flow matrix
# output:
#  backward total flow vector
  '''
  return numpy.array((numpy.matrix(numpy.ones(flowmat.shape[0]))*flowmat).tolist());

def cmbfMat(flowmat):
  '''
# Calculate forward combination criteria matrix of a given flow matrix
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  flow matrix
# output:
#  forward combination criteria matrix
  '''
  tmpmat=numpy.matrix(2.*numpy.array(simfMat(flowmat)+1.)/(numpy.array(tflfVec(flowmat).transpose()*numpy.ones(flowmat.shape[1]))+1.)/(numpy.array(flowmat+flowmat.transpose())+1.));
  return numpy.matrix(numpy.array(tmpmat)*numpy.array(tmpmat.transpose()));

def cmbbMat(flowmat):
  '''
# Calculate backward combination criteria matrix of a given flow matrix
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  flow matrix
# output:
#  backward combination criteria matrix
  '''
  tmpmat=numpy.matrix(2.*numpy.array(simbMat(flowmat)+1.)/(numpy.array(numpy.ones(flowmat.shape[0]).transpose()*tflbVec(flowmat))+1.)/(numpy.array(flowmat+flowmat.transpose())+1.));
  return numpy.matrix(numpy.array(tmpmat)*numpy.array(tmpmat.transpose()));

def numVec(clusters):
  '''
# Calculate number-vector
# number-vector: number of papers the cluster contains
# input:
#  clusters list
# output:
#  numvec
  '''
  dimc=len(clusters);
  numvec=[len(clusters[c]) for c in range(0,dimc)];
  return numvec;

def simplify(lastlist,newclusters,lasthist=None):
  '''
# Simplify the recursive history by once
# Generate the tree history from recursive history (if lasthist is not None)
# by Sun Smallwhite <niasw@pku.edu.cn>(https://github.com/niasw)

# input:
#  the latest clusters list, new clusters dump list, the latest history node (assign None to turn off)
# output:
#  new clusters list, the new history node (if lasthist is not None), the modified latest history node (if lasthist is not None)
  '''
  if (lasthist is None):
    if (lastlist is None):
      for itm in range(0,len(newclusters)):
        newclusters[itm].sort();
      return newclusters;
    else:
      curlist=[];
      for itm in newclusters:
        tmp=[];
        for itp in itm:
          tmp.extend(lastlist[itp]);
        tmp.sort();
        curlist.append(tmp);
      return curlist;
  else:
    if (lastlist is None):
      for itm in range(0,len(newclusters)):
        newclusters[itm].sort();
      return newclusters,[{"id":(str(len(newclusters))+'-'+str(cit)),"prt":None,"chd":None,"clus":newclusters[cit]} for cit in range(0,len(newclusters))],[];
    else:
      curlist=[];
      curhist=[];
      for cit in range(0,len(newclusters)):
        tmp=[];
        chdlist=[];
        name=str(len(newclusters))+'-'+str(cit);
        for itp in newclusters[cit]:
          tmp.extend(lastlist[itp]);
          chdlist.append(str(len(lasthist))+'-'+str(itp));
          lasthist[itp]["prt"]=name;
        tmp.sort();
        curlist.append(tmp);
        curhist.append({"id":name,"prt":None,"chd":chdlist});
      return curlist,curhist,lasthist;
  return None;

def calcXvector(clusters,dimension,numvec=None):
  '''
# for a paper p:
#  x(p)=sqrt(|#p|)/sqrt(|#c(p)|)
# compatible without numvec
# input:
#  clusters: the division result
#  numvec: paper number of initial clusters
#  dimension: number of initial clusters
  '''
  xvec=numpy.zeros(dimension)+1.;
  for c in clusters:
   if (numvec is None):
    for p in c:
     xvec[p]=1./numpy.sqrt(numpy.float64(len(c)));
   else:
    csum=numpy.array(numvec)[c].sum();
    for p in c:
     xvec[p]=numpy.sqrt(numpy.float64(numvec[p])/numpy.float64(csum));
  return xvec;

